package com.example.state;

import org.apache.flink.api.common.functions.RichFlatMapFunction;
import org.apache.flink.api.common.state.ListState;
import org.apache.flink.api.common.state.ListStateDescriptor;
import org.apache.flink.api.common.state.ValueState;
import org.apache.flink.api.common.state.ValueStateDescriptor;
import org.apache.flink.api.common.typeinfo.TypeHint;
import org.apache.flink.api.common.typeinfo.TypeInformation;
import org.apache.flink.api.java.tuple.Tuple2;
import org.apache.flink.configuration.Configuration;
import org.apache.flink.runtime.state.FunctionInitializationContext;
import org.apache.flink.runtime.state.FunctionSnapshotContext;
import org.apache.flink.streaming.api.checkpoint.CheckpointedFunction;
import org.apache.flink.streaming.api.datastream.DataStreamSource;
import org.apache.flink.streaming.api.datastream.KeyedStream;
import org.apache.flink.streaming.api.datastream.SingleOutputStreamOperator;
import org.apache.flink.streaming.api.environment.StreamExecutionEnvironment;
import org.apache.flink.streaming.api.functions.sink.SinkFunction;
import org.apache.flink.util.Collector;

import java.util.ArrayList;
import java.util.List;

/**
 * @author lbq
 * @version 1.0
 * @date 2022/3/27 14:52
 */
public class ListCheakpointedDemo {
    public static void main(String[] args) throws Exception {
        StreamExecutionEnvironment env = StreamExecutionEnvironment.getExecutionEnvironment();
        // env.enableCheckpointing(2);
        final DataStreamSource<Tuple2<Integer, Integer>> data = env.fromElements(
                Tuple2.of(1, 3),
                Tuple2.of(1, 5),
                Tuple2.of(1, 7),
                Tuple2.of(1, 4),
                Tuple2.of(1, 2)
        );

        final KeyedStream<Tuple2<Integer, Integer>, Integer> keyedStream = data.keyBy(value -> value.f0);

        final SingleOutputStreamOperator<Tuple2<Integer, Integer>> process = keyedStream.flatMap(new RichFlatMapFunction<Tuple2<Integer, Integer>, Tuple2<Integer, Integer>>() {

            private ValueState<Tuple2<Integer, Integer>> sum;

            @Override
            public void flatMap(Tuple2<Integer, Integer> value, Collector<Tuple2<Integer, Integer>> out) throws Exception {
                final Tuple2<Integer, Integer> currentState = sum.value();
                currentState.f0 += 1;
                currentState.f1 += value.f1;
                System.out.println("currentState = " + currentState);
                sum.update(currentState);

                if (currentState.f0 >= 2) {
                    out.collect(new Tuple2<>(value.f0, currentState.f1 / currentState.f0));
                    // sum.clear();
                }
            }

            @Override
            public void open(Configuration parameters) throws Exception {
                final ValueStateDescriptor<Tuple2<Integer, Integer>> descriptor = new ValueStateDescriptor<>("average", TypeInformation.of(new TypeHint<Tuple2<Integer, Integer>>() {
                }), Tuple2.of(0, 0));
                sum = getRuntimeContext().getState(descriptor);
            }
        });

        process.addSink(new BufferSink(2));

        env.execute();
    }

    public static class BufferSink implements SinkFunction<Tuple2<Integer, Integer>>, CheckpointedFunction {

        private ListState<Tuple2<Integer, Integer>> checkPointedState;
        private final List<Tuple2<Integer, Integer>> bufferedElements;
        private final int threads;

        public BufferSink(int threads) {
            bufferedElements = new ArrayList<>();
            this.threads = threads;
        }

        @Override
        public void invoke(Tuple2<Integer, Integer> value, Context context) throws Exception {
            bufferedElements.add(value);
            if (bufferedElements.size() == threads) {
                for (Tuple2<Integer, Integer> element : bufferedElements) {
                    System.out.println("...out:" + element);
                }
                bufferedElements.clear();
            }
        }

        @Override
        public void snapshotState(FunctionSnapshotContext context) throws Exception {
            checkPointedState.clear();
            for (Tuple2<Integer, Integer> tuple2 : bufferedElements) {
                checkPointedState.add(tuple2);
            }
        }

        @Override
        public void initializeState(FunctionInitializationContext context) throws Exception {
            final ListStateDescriptor<Tuple2<Integer, Integer>> descriptor = new ListStateDescriptor<>("buffered", TypeInformation.of(new TypeHint<Tuple2<Integer, Integer>>() {
            }));
            checkPointedState = context.getOperatorStateStore().getListState(descriptor);
            if (context.isRestored()) {
                for (Tuple2<Integer, Integer> tuple2 : checkPointedState.get()) {
                    bufferedElements.add(tuple2);
                }
            }
        }
    }
}
